import numpy as np
from math import sin, cos, radians, sqrt, degrees, atan
from typing import Tuple

# --- I. 物理常数和弹丸参数 ---
RHO0 = 1.225         # 海平面空气密度 (kg/m^3)
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 绝热垂直温度梯度 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 干燥空气比气体常数 (J/kg·K)
P_TERM = (G / (R * L)) - 1 # 用于计算密度比的指数
C_SOUND = 338.0      # 海平面声速 (m/s)

M = 5.0              # 弹丸质量 (kg)
A = 0.0095           # 弹丸截面积 (m^2)
V_WIND_X = 0.0       # 纵向风速度 (固定为零)
H0 = 0.0             # 初始发射和落地海拔高度 (固定为海平面)

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) # 马赫数查找点
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) # 阻力系数

# --- II. 动态物理模型和RK4函数 (2D简化版) ---

def get_drag_coefficient(v_rel):
    """根据相对速度计算空气阻力系数 (Cd)。"""
    mach_number = v_rel / C_SOUND
    return np.interp(mach_number, MA_POINTS, CD_POINTS)

def get_air_density(z):
    """根据高度 Z 计算空气密度 (假设温度随高度变化)。"""
    h = max(0.0, z) 
    # 此处 H0 始终为 0，但 Z 在飞行中会变化，所以仍需计算空气密度
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0: return 0.0 
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

def derivatives_of_state_2d(S: np.ndarray, params: Tuple[float, float, float]) -> np.ndarray:
    """计算 2D 状态向量 S = [x, z, vx, vz] 在时间上的微分。"""
    x, z, vx, vz = S
    M, A, G = params 
    
    # 1. 计算相对速度 (纵向风 V_WIND_X = 0.0)
    vx_rel = vx - V_WIND_X 
    vz_rel = vz
    v_rel = sqrt(vx_rel**2 + vz_rel**2) # 相对速度的模

    if v_rel < 1e-6:
        return np.array([vx, vz, 0.0, -G])

    rho = get_air_density(z) # 密度随飞行高度 Z 变化
    Cd = get_drag_coefficient(v_rel)

    # 2. 空气阻力加速度
    K = -(Cd * rho * A * v_rel) / (2.0 * M)
    
    dvx_drag = K * vx_rel
    dvz_drag = K * vz_rel

    # 3. 合并加速度
    dvx_dt = dvx_drag
    dvz_dt = dvz_drag - G

    return np.array([vx, vz, dvx_dt, dvz_dt])

def rk4_step_2d(S: np.ndarray, dt: float, params: Tuple[float, float, float]) -> np.ndarray:
    """四阶龙格-库塔方法进行一步 2D 积分。"""
    f = lambda S: derivatives_of_state_2d(S, params)
    
    K1 = dt * f(S)
    K2 = dt * f(S + K1 / 2.0)
    K3 = dt * f(S + K2 / 2.0)
    K4 = dt * f(S + K3)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

def simulate_and_get_impact(v0: float, theta_deg: float, dt: float) -> Tuple[float, float]:
    """
    模拟 2D 轨迹并返回最终落点 X 坐标 (射程) 和落地角度 (gamma)。
    """
    params = (M, A, G)
    theta_rad = radians(theta_deg)
    
    # 2D 状态向量: [x, z, vx, vz]. 初始高度 Z = H0 = 0.0
    S = np.array([
        0.0, H0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    # 假设炮弹在初始高度 H0 = 0 落地
    while S[1] >= H0 - 0.01:
        S_new = rk4_step_2d(S, dt, params) 
        
        # 冲击检查 (Z 穿过 H0 平面)
        if S_new[1] < H0 and S[1] >= H0:
            # 1. 线性插值计算冲击点 X 坐标 (射程)
            t_impact_ratio = (S[1] - H0) / (S[1] - S_new[1])
            R_final = S[0] + (S_new[0] - S[0]) * t_impact_ratio

            # 2. 估计冲击时的速度分量 (也进行插值以提高精度)
            vx_final = S[2] + (S_new[2] - S[2]) * t_impact_ratio
            vz_final = S[3] + (S_new[3] - S[3]) * t_impact_ratio
            
            # 3. 计算落地角度 (gamma_impact) - 速度矢量与水平方向的夹角 (0到90度)
            if abs(vx_final) < 1e-6:
                 gamma_impact_deg = 90.0
            else:
                 # vz_final 在落地时为负值，atan只取比值的绝对值
                 gamma_impact_deg = degrees(atan(abs(vz_final) / abs(vx_final)))
                 
            return R_final, gamma_impact_deg
            
        S = S_new
        
    return S[0], np.nan # 落地角度找不到，或落在了很远的地方

# --- III. 搜索函数 (计算 V0 和 theta) ---

def binary_search_v0_for_range(target_range: float, theta_deg: float, dt: float, V_low=10.0, V_high=2000.0, tolerance=1.0) -> float:
    """查找固定角度、目标距离所需的 V0。"""
    
    h0 = H0 # 使用固定的海平面高度
    
    for _ in range(50):
        V_mid = (V_low + V_high) / 2.0
        # 只需要射程，所以调用 simulate_and_get_impact 并忽略 gamma
        R_mid, _ = simulate_and_get_impact(V_mid, theta_deg, dt)
        
        if abs(R_mid - target_range) < tolerance:
            return V_mid
        
        if R_mid < target_range:
            V_low = V_mid 
        else:
            V_high = V_mid 
            
    return np.nan

def search_v0_theta_for_constraints(target_range: float, target_gamma_deg: float, dt: float) -> Tuple[float, float]:
    """
    搜索满足目标射程和目标落地角度的 (V0, Theta) 组合。
    """
    
    # 角度搜索范围和步长
    THETA_MIN = 10.0
    THETA_MAX = 70.0
    THETA_STEP = 0.5  

    best_v0 = np.nan
    best_theta = np.nan
    min_gamma_error = float('inf')
    
    # 外层循环: 遍历可能的发射角 theta
    for theta_deg in np.arange(THETA_MIN, THETA_MAX + THETA_STEP, THETA_STEP):
        
        # 内层搜索: 找到 V0，使得 (V0, theta) 组合能击中 target_range
        V0_required = binary_search_v0_for_range(
            target_range=target_range, 
            theta_deg=theta_deg, 
            dt=dt
        )
        
        if np.isnan(V0_required):
            continue

        # 使用找到的 (V0, theta) 组合进行精确模拟，获取落地角度
        R_actual, gamma_impact_deg = simulate_and_get_impact(V0_required, theta_deg, dt)
        
        if np.isnan(gamma_impact_deg):
            continue

        # 计算角度误差
        gamma_error = abs(gamma_impact_deg - target_gamma_deg)

        # 找到角度最接近目标的 (V0, Theta)
        if gamma_error < min_gamma_error:
            min_gamma_error = gamma_error
            best_v0 = V0_required
            best_theta = theta_deg
            
        # 优化：如果当前发射角 theta 带来的落地角 gamma 已经远大于目标，停止搜索。
        # 对于固定射程，增加 theta 通常导致更陡峭的落地角。
        if gamma_impact_deg > target_gamma_deg + 2.0: 
            break

    # 如果最小误差仍大于可接受的范围 (例如 1.0 度)，则认为找不到解
    if min_gamma_error > 1.0: 
        return np.nan, np.nan
        
    return best_v0, best_theta

# --- IV. 主运行和表格输出 ---

def generate_range_impact_angle_table():
    
    # --- 固定参数 ---
    DT = 0.1             # RK4 步长 (s)

    # --- 表格参数 ---
    R_MIN = 900.0
    R_MAX = 1600.0
    R_STEP = 100.0        # 目标距离步长
    GAMMA_MIN = 10.0
    GAMMA_MAX = 80.0
    GAMMA_STEP = 10.0     # 目标落地角度步长 (50, 60, 70, 80度)
    
    R_TARGETS = np.arange(R_MIN, R_MAX + R_STEP, R_STEP)
    GAMMA_TARGETS = np.arange(GAMMA_MIN, GAMMA_MAX + GAMMA_STEP, GAMMA_STEP)

    if len(R_TARGETS) < 2 or len(GAMMA_TARGETS) < 2:
        print("❌ 错误: 步长设置不合理，无法生成表格。")
        return

    # 存储结果的二维数组 (存储 (V0, Theta) 元组)
    results_table = np.empty((len(GAMMA_TARGETS), len(R_TARGETS)), dtype=object)

    print("--- 🎯 正在进行 2D 射程和落地角度双重约束下的发射参数搜索 ---")
            
    # 填充表格
    for j, R_target in enumerate(R_TARGETS):
        print(f"\n[目标射程 R={R_target:.0f}m] 正在搜索所需参数...")
        
        for i, Gamma_target in enumerate(GAMMA_TARGETS):
            
            # 搜索满足 R_target 和 Gamma_target 的 (V0, Theta) 组合
            V0_opt, Theta_opt = search_v0_theta_for_constraints(
                target_range=R_target, 
                target_gamma_deg=Gamma_target, 
                dt=DT
            )
            
            # 存储 (V0, Theta) 元组
            results_table[i, j] = (V0_opt, Theta_opt)
            
            # 打印调试信息
            V0_str = f"{V0_opt:.2f}" if not np.isnan(V0_opt) else "N/A"
            Theta_str = f"{Theta_opt:.1f}" if not np.isnan(Theta_opt) else "N/A"
            
            print(f"  R:{R_target:4.0f}m, γ:{Gamma_target:2.0f}° => V₀:{V0_str}m/s, θ:{Theta_str}°")


    # --- 结果表格输出 ---
    print("\n" + "="*120)
    print("      🎯 满足给定射程和落地角度约束的发射参数 (V₀ (m/s), θ (度)) 查表 🎯")
    print("      [注: V₀ 为初速度，θ 为出射角。仅考虑重力与随速度变化的空气阻力]")
    print("="*120)

    # 格式化表头 (X轴：目标直线距离 R (m))
    header_cols = [f"{r:^15.0f}" for r in R_TARGETS]
    header = " | " + " | ".join(header_cols) + " |"
    print(" 落地角度 γ (度) \\ 目标距离 R (m)")
    print("-" * len(header))
    print(header)
    print("-" * len(header))
    
    # 格式化数据行 (Y轴：目标落地角度 Gamma (度))
    for i, Gamma in enumerate(GAMMA_TARGETS):
        row = []
        for j in range(len(R_TARGETS)):
            data = results_table[i, j]
            if isinstance(data, tuple) and not np.isnan(data[0]):
                V0, Theta = data
                # 格式化为 V0.xx / T.x
                formatted_pair = f"{V0:.1f} / {Theta:.1f}"
                row.append(f"{formatted_pair:^15}")
            else:
                row.append(f"{'N/A':^15}")
        
        print(f" | {Gamma:^10.0f} | " + " | ".join(row) + " |")

    print("="*120)
    print("💡 查表指引:")
    print("   - V₀ / θ: 满足该射程 (R) 和该落地角度 (γ) 所需的最小初速度 (m/s) 及其对应的出射角 (度)。")
    print("   - 随着所需落地角度 (γ) 变大，通常需要更大的 V₀ 或更高的 θ 来实现更陡峭的落角。")
    
if __name__ == "__main__":
    generate_range_impact_angle_table()